package com.nexus.ai.vectordb;

import org.springframework.ai.document.Document;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.ai.vectorstore.filter.Filter;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import java.util.List;
import java.util.Map;

/**
 * 向量数据库操作类
 *
 * @Date 2025/9/28 19:13
 * @Author luzhengning
 **/
@Service
public class NexusVectordbService {

    @Autowired
    private VectorStore vectorStore;

    /**
     * 保存文档到向量数据库
     *
     * @Date 2025/9/28 19:49
     * @Author luzhengning
     **/
    public void save() {
        //创建文档
        Document doc1 = new Document("doc-001", "我的电台型号有八重洲FT-818nd,FT-4x和UV5rh",
                Map.of(
                        "source", "电台库存",  //来源
                        "category", "电台型号", //分类
                        "created_at", System.currentTimeMillis()));

        Document doc2 = new Document("doc-002", "西安有肉夹馍，兰州有牛肉面，新疆有大盘鸡",
                Map.of(
                        "source", "美食中国",//来源
                        "category", "地方特色",  //分类
                        "created_at", System.currentTimeMillis()));

        //集合组织多份文档
        List<Document> docs = List.of(doc1, doc2);
        //添加到向量数据库
        vectorStore.add(docs);
    }

    /**
     * 查询文档
     *
     * @Date 2025/9/28 20:29
     * @Author luzhengning
     **/
    public List<Document> search(String queryMsg) {
//        //执行相似性查询，相似性从高到低排序
//        List<Document> documents = vectorStore.similaritySearch(queryMsg);
//        //使用相似性得分阈值过滤结果
//        SearchRequest build = SearchRequest.builder().query(queryMsg)
//                .similarityThreshold(0.4) //相似性阈值
//                .build();
//        List<Document> documents = vectorStore.similaritySearch(build);

//        //找出最相似的前几个文档
//        SearchRequest build = SearchRequest.builder().query(queryMsg)
//                .topK(1)    //最相似的前几个文档
//                .build();
//        List<Document> documents = vectorStore.similaritySearch(build);

//        //元数据过滤查询
//        SearchRequest build = SearchRequest.builder().query(queryMsg)
//                .topK(1)    //最相似的前几个文档
//                .filterExpression("source=='美食中国'")
//                .build();
//        List<Document> documents = vectorStore.similaritySearch(build);

        //条件表达式查询/元数据过滤
        Filter.Expression filter = new Filter.Expression(
                Filter.ExpressionType.LT,//大于
                new Filter.Key("created_at"),
                new Filter.Value(System.currentTimeMillis())
        );
        SearchRequest build = SearchRequest.builder().query(queryMsg)
                .topK(1)    //最相似的前几个文档
                .filterExpression(filter)
                .build();
        List<Document> documents = vectorStore.similaritySearch(build);
        return documents;
    }

    public void delete() {
        //根据文档ID删除
        //vectorStore.delete(List.of("doc-001","doc-002"));
        //根据元数据删除
        //vectorStore.delete("source == '美食中国'");
        //根据元数据删除
        vectorStore.delete(new Filter.Expression(
                Filter.ExpressionType.EQ,
                new Filter.Key("source"),
                new Filter.Value("美食中国")
        ));
    }
}
